#!/usr/bin/env python3
"""
更新数据配置文件，移除D90类别
"""

import os
from pathlib import Path

def update_data_configs():
    """更新所有数据配置文件，移除D90类别"""
    
    # 更新主配置文件
    main_config = Path("N-RDD2024/data.yaml")
    if main_config.exists():
        new_content = """train: /root/autodl-tmp/yolo11n_ljy/ultralytics/N-RDD2024/merged_dataset/images/train
val: /root/autodl-tmp/yolo11n_ljy/ultralytics/N-RDD2024/merged_dataset/images/val

nc: 9
names: ['D00', 'D10', 'D20', 'D30', 'D40', 'D50', 'D60', 'D70', 'D80']
"""
        with open(main_config, 'w', encoding='utf-8') as f:
            f.write(new_content)
        print(f"已更新: {main_config}")
    
    # 更新各国家的配置文件
    base_path = Path("N-RDD2024/Training and Validation Dataset")
    
    for country_dir in base_path.iterdir():
        if not country_dir.is_dir() or not country_dir.name.endswith('_txt'):
            continue
            
        config_file = country_dir / "data.yaml"
        if config_file.exists():
            # 读取现有配置
            with open(config_file, 'r', encoding='utf-8') as f:
                content = f.read()
            
            # 更新类别数和名称
            lines = content.split('\n')
            new_lines = []
            
            for line in lines:
                if line.strip().startswith('nc:'):
                    new_lines.append('nc: 9')
                elif 'names:' in line or any(f"'{cls}'" in line for cls in ['D00', 'D10', 'D20', 'D30', 'D40', 'D50', 'D60', 'D70', 'D80', 'D90']):
                    # 跳过旧的names部分，稍后添加新的
                    continue
                else:
                    new_lines.append(line)
            
            # 添加新的names配置
            new_lines.append('')
            new_lines.append('names:')
            for i, cls in enumerate(['D00', 'D10', 'D20', 'D30', 'D40', 'D50', 'D60', 'D70', 'D80']):
                new_lines.append(f'  {i}: {cls}')
            
            # 写回文件
            with open(config_file, 'w', encoding='utf-8') as f:
                f.write('\n'.join(new_lines))
            
            print(f"已更新: {config_file}")
    
    # 更新根目录的data.yaml
    root_config = Path("data.yaml")
    if root_config.exists():
        new_content = """# RDD2022 数据集配置文件
# 路径配置
train: RDD2022/images/train
val: RDD2022/images/val
test:  # 如果需要测试集，请添加路径
# 类别数量
nc: 9  # 移除D90后的类别数量

# 类别名称列表
names:
  - D00  # 纵向裂缝
  - D10  # 横向裂缝
  - D20  # 网状裂缝
  - D30  # 
  - D40  # 修补痕迹
  - D50  # 
  - D60  # 
  - D70  # 
  - D80  # 

# 可选参数
# download: https://github.com/sekilab/RoadDamageDetector  # 数据集来源"""
        
        with open(root_config, 'w', encoding='utf-8') as f:
            f.write(new_content)
        print(f"已更新: {root_config}")
    
    print("\n配置文件更新完成!")

if __name__ == "__main__":
    update_data_configs()